package com.jxd.dianping.recommend;

import lombok.Data;
import org.apache.commons.lang3.StringUtils;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.function.ForeachPartitionFunction;
import org.apache.spark.api.java.function.Function;
import org.apache.spark.ml.recommendation.ALSModel;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.catalyst.expressions.GenericRowWithSchema;

import java.io.Serializable;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.util.*;

// 文件读取有问题 问题出在windows环境中没有运行hadoop组件
// 在hadoop/bin中加入winutil和hadoop.dll
// 并把hadoop的文件夹位置写入系统环境变量
public class AlsRecallPredict {


    public static void main(String[] args) {
        // 初始化spark的运行环境
        SparkSession sparkSession = SparkSession.builder().master("local").appName("DianpingApp").getOrCreate();
        // 加载训练后的模型
        ALSModel alsModel = ALSModel.load("file:///F:\\IdeaProjects\\self-learn\\dianping\\src\\main\\resources\\model\\alsmodel");
        // 读取数据
        JavaRDD<String> csvFile = sparkSession.read().textFile("file:///F:\\IdeaProjects\\self-learn\\dianping\\src\\main\\resources\\data\\behavior.csv").toJavaRDD();
        // 将list String 转为 list Rating
        JavaRDD<AlsRecallTrain.Rating> ratingJavaRDD = csvFile.map(new Function<String, AlsRecallTrain.Rating>() {
            @Override
            public AlsRecallTrain.Rating call(String s){
                return AlsRecallTrain.Rating.parseRating(s);
            }
        });
        // 以rating为字段的数据表
        Dataset<Row> ratingDataFrame = sparkSession.createDataFrame(ratingJavaRDD, AlsRecallTrain.Rating.class);
        // 给5个用户做离线结果召回
        Dataset<Row> users = ratingDataFrame.select(alsModel.getUserCol()).distinct().limit(5);
        Dataset<Row> userRecs = alsModel.recommendForUserSubset(users, 20);
        // list分片
        userRecs.foreachPartition(new ForeachPartitionFunction<Row>() {
            @Override
            public void call(Iterator<Row> iterator) throws Exception {
                // 新建数据库连接
                Connection connection = DriverManager.getConnection("jdbc:mysql://127.0.0.1:3306/dianpingdb?user=root&password=123456&useUnicode=true&characterEncoding=UTF-8&useSSL=false");
                PreparedStatement preparedStatement = connection.prepareStatement("insert into recommend(id, recommend) values(?,?)");
                List<Map<String,Object>> data = new ArrayList<>();
                // 遍历sublist
                iterator.forEachRemaining(action->{
                    int userId = action.getInt(0);
                    List<GenericRowWithSchema> list = action.getList(1);
                    List<Integer> shopList = new ArrayList<>();
                    list.forEach(row->{
                        Integer shopId = row.getInt(0);
                        shopList.add(shopId);
                    });
                    String recommendData = StringUtils.join(shopList,",");
                    Map<String, Object> map = new HashMap<>();
                    map.put("userId", userId);
                    map.put("recommend", recommendData);
                    data.add(map);
                });
                //
                data.forEach(stringObjectMap -> {
                    try {
                        preparedStatement.setInt(1, (Integer) stringObjectMap.get("userId"));
                        preparedStatement.setString(2, (String)stringObjectMap.get("recommend"));
                        preparedStatement.addBatch();
                    } catch (SQLException throwables) {
                        throwables.printStackTrace();
                    }
                    try {
                        preparedStatement.executeBatch();
                    } catch (SQLException throwables) {

                    }
                    try {
                        preparedStatement.close();
                    } catch (SQLException throwables) {
                        throwables.printStackTrace();
                    }
                });
            }
        });
    }

    @Data
    public static class Rating implements Serializable {
        private int userId;
        private int shopId;
        private int rating;

        public Rating(int userId, int shopId, int rating) {
            this.userId = userId;
            this.shopId = shopId;
            this.rating = rating;
        }

        public static AlsRecallTrain.Rating parseRating(String str) {
            str = str.replace("\"", "");
            String[] strArr = str.split(",");
            int userId = Integer.parseInt(strArr[0]);
            int shopId = Integer.parseInt(strArr[1]);
            int rating = Integer.parseInt(strArr[2]);
            return new AlsRecallTrain.Rating(userId, shopId, rating);
        }
    }
}
